package com.haha.sum;

import net.sourceforge.tess4j.ITesseract;
import net.sourceforge.tess4j.Tesseract;
import net.sourceforge.tess4j.TesseractException;
import net.sourceforge.tess4j.Word;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * @author shenlinnan
 * @date 2025/7/27 20:56
 */
public class WoWGroupListScanner {

    // 配置参数
    private static final String IMAGE_PATH = "I:\\my-workspace\\haha-maven\\screenshots\\screenshot_20250727-171422.png";
    // 修正语言数据路径 - 指向包含 eng.traineddata 的目录
    private static final String TESSDATA_PATH = "I:\\my-workspace\\haha-maven\\src\\main\\resources";
    private static final String LANGUAGE = "eng"; // 英文识别

    public static void main(String[] args) {
        try {
            // 1. 加载图像
            BufferedImage image = ImageIO.read(new File(IMAGE_PATH));
            System.out.println("图像尺寸: " + image.getWidth() + "×" + image.getHeight());

            ITesseract tesseract = new Tesseract();
            tesseract.setDatapath(TESSDATA_PATH);
            tesseract.setLanguage(LANGUAGE);
            List<Word> words = tesseract.getWords(image, ITesseract.RenderedFormat.TEXT.ordinal());

            // 2. 预处理图像（增强对比度）
            BufferedImage processedImage = preprocessImage(image);

            // 3. 使用Tess4J进行OCR识别
            List<TextInfo> textInfos = recognizeText(processedImage);

            // 4. 输出结果
            System.out.println("\n=== 魔兽世界组队列表识别结果 ===");
            for (TextInfo info : textInfos) {
                System.out.printf("角色: %-15s | 颜色: RGB(%3d, %3d, %3d)%n",
                        info.text,
                        info.color.getRed(),
                        info.color.getGreen(),
                        info.color.getBlue());
            }

            // 5. 生成带标注的图像
            createAnnotatedImage(image, textInfos);
            System.out.println("\n已创建带标注的图像: wow_group_list_annotated.png");

        } catch (IOException | TesseractException e) {
            System.err.println("处理错误: " + e.getMessage());
            e.printStackTrace();
        }
    }

    /**
     * 图像预处理 - 增强对比度，提高OCR识别率
     */
    private static BufferedImage preprocessImage(BufferedImage original) {
        BufferedImage processed = new BufferedImage(
                original.getWidth(),
                original.getHeight(),
                BufferedImage.TYPE_INT_RGB
        );

        // 增强对比度
        for (int y = 0; y < original.getHeight(); y++) {
            for (int x = 0; x < original.getWidth(); x++) {
                Color color = new Color(original.getRGB(x, y));

                // 增强颜色对比度
                int r = enhanceContrast(color.getRed());
                int g = enhanceContrast(color.getGreen());
                int b = enhanceContrast(color.getBlue());

                processed.setRGB(x, y, new Color(r, g, b).getRGB());
            }
        }

        return processed;
    }

    /**
     * 增强颜色对比度
     */
    private static int enhanceContrast(int value) {
        // 非线性增强 - 使亮色更亮，暗色更暗
        if (value > 128) {
            return Math.min(255, value + (value - 128) / 2);
        } else {
            return Math.max(0, value - (128 - value) / 2);
        }
    }

    /**
     * 识别文本和颜色
     */
    private static List<TextInfo> recognizeText(BufferedImage image) throws TesseractException {
        ITesseract tesseract = new Tesseract();
        tesseract.setDatapath(TESSDATA_PATH);
        tesseract.setLanguage(LANGUAGE);

        // 获取识别结果（带位置信息）
        List<Word> words = tesseract.getWords(image, ITesseract.RenderedFormat.TEXT.ordinal());

        // 按行分组（假设每行文字高度为20像素）
        List<List<Word>> lines = groupWordsIntoLines(words, 20);

        // 处理每行文字
        List<TextInfo> results = new ArrayList<>();
        for (List<Word> line : lines) {
            if (!line.isEmpty()) {
                // 合并行内所有单词
                String text = combineWords(line);

                // 计算整行的平均颜色
                Color color = calculateLineColor(image, line);

                results.add(new TextInfo(text, color));
            }
        }

        return results;
    }

    /**
     * 将单词按行分组
     */
    private static List<List<Word>> groupWordsIntoLines(List<Word> words, int lineHeight) {
        List<List<Word>> lines = new ArrayList<>();

        // 如果没有单词，返回空列表
        if (words.isEmpty()) return lines;

        // 按Y坐标排序
        words.sort((w1, w2) -> Integer.compare(w1.getBoundingBox().y, w2.getBoundingBox().y));

        // 分组
        List<Word> currentLine = new ArrayList<>();
        int currentY = words.get(0).getBoundingBox().y;

        for (Word word : words) {
            int wordY = word.getBoundingBox().y;

            // 如果单词在同一行（在可接受的Y坐标范围内）
            if (Math.abs(wordY - currentY) <= lineHeight) {
                currentLine.add(word);
            } else {
                // 新行开始
                if (!currentLine.isEmpty()) {
                    lines.add(currentLine);
                }
                currentLine = new ArrayList<>();
                currentLine.add(word);
                currentY = wordY;
            }
        }

        // 添加最后一行
        if (!currentLine.isEmpty()) {
            lines.add(currentLine);
        }

        return lines;
    }

    /**
     * 合并行内所有单词
     */
    private static String combineWords(List<Word> words) {
        // 按X坐标排序
        words.sort((w1, w2) -> Integer.compare(w1.getBoundingBox().x, w2.getBoundingBox().x));

        StringBuilder sb = new StringBuilder();
        for (Word word : words) {
            if (sb.length() > 0) {
                sb.append(" ");
            }
            sb.append(word.getText());
        }

        return sb.toString();
    }

    /**
     * 计算整行的平均颜色
     */
    private static Color calculateLineColor(BufferedImage image, List<Word> words) {
        int redSum = 0, greenSum = 0, blueSum = 0;
        int pixelCount = 0;

        // 计算整行的边界
        Rectangle lineRect = calculateLineBoundingBox(words);

        // 遍历行内的每个单词
        for (Word word : words) {
            Rectangle wordRect = word.getBoundingBox();

            // 计算单词的中心区域（避免边缘干扰）
            int centerX = wordRect.x + wordRect.width / 2;
            int centerY = wordRect.y + wordRect.height / 2;
            int sampleWidth = Math.max(5, wordRect.width / 2);
            int sampleHeight = Math.max(5, wordRect.height / 2);

            // 在单词中心区域采样颜色
            for (int y = centerY - sampleHeight/2; y < centerY + sampleHeight/2; y++) {
                for (int x = centerX - sampleWidth/2; x < centerX + sampleWidth/2; x++) {
                    if (x >= 0 && y >= 0 && x < image.getWidth() && y < image.getHeight()) {
                        Color color = new Color(image.getRGB(x, y));

                        // 跳过黑色背景
                        if (!isBackground(color)) {
                            redSum += color.getRed();
                            greenSum += color.getGreen();
                            blueSum += color.getBlue();
                            pixelCount++;
                        }
                    }
                }
            }
        }

        // 如果没有有效像素，返回黑色
        if (pixelCount == 0) {
            return Color.BLACK;
        }

        return new Color(
                redSum / pixelCount,
                greenSum / pixelCount,
                blueSum / pixelCount
        );
    }

    /**
     * 计算整行的边界框
     */
    private static Rectangle calculateLineBoundingBox(List<Word> words) {
        if (words.isEmpty()) return new Rectangle();

        int minX = Integer.MAX_VALUE;
        int minY = Integer.MAX_VALUE;
        int maxX = 0;
        int maxY = 0;

        for (Word word : words) {
            Rectangle rect = word.getBoundingBox();
            minX = Math.min(minX, rect.x);
            minY = Math.min(minY, rect.y);
            maxX = Math.max(maxX, rect.x + rect.width);
            maxY = Math.max(maxY, rect.y + rect.height);
        }

        return new Rectangle(minX, minY, maxX - minX, maxY - minY);
    }

    /**
     * 判断是否是背景色（黑色）
     */
    private static boolean isBackground(Color color) {
        // 检测黑色背景（RGB值都很低）
        return color.getRed() < 50 && color.getGreen() < 50 && color.getBlue() < 50;
    }

    /**
     * 创建带标注的图像
     */
    private static void createAnnotatedImage(BufferedImage original, List<TextInfo> textInfos) {
        // 创建副本图像
        BufferedImage annotated = new BufferedImage(
                original.getWidth(),
                original.getHeight(),
                BufferedImage.TYPE_INT_RGB
        );

        // 绘制原始图像
        Graphics2D g2d = annotated.createGraphics();
        g2d.drawImage(original, 0, 0, null);

        // 设置标注样式
        Font font = new Font("Arial", Font.BOLD, 14);
        g2d.setFont(font);

        // 为每个文本添加标注
        int yOffset = 20;
        for (TextInfo info : textInfos) {
            // 绘制颜色框
            g2d.setColor(info.color);
            g2d.fillRect(10, yOffset, 20, 20);

            // 绘制文本
            g2d.setColor(Color.WHITE);
            g2d.drawString(info.text, 40, yOffset + 15);

            // 绘制RGB值
            String rgb = String.format("RGB(%d, %d, %d)",
                    info.color.getRed(), info.color.getGreen(), info.color.getBlue());
            g2d.drawString(rgb, 200, yOffset + 15);

            yOffset += 30;
        }

        g2d.dispose();

        // 保存标注后的图像
        try {
            ImageIO.write(annotated, "png", new File("wow_group_list_annotated.png"));
        } catch (IOException e) {
            System.err.println("保存标注图像失败: " + e.getMessage());
        }
    }

    /**
     * 封装文本信息
     */
    private static class TextInfo {
        String text;
        Color color;

        public TextInfo(String text, Color color) {
            this.text = text;
            this.color = color;
        }
    }

}
